#include "ompl/base/goals/GoalState.h"
#include "ompl/base/goals/GoalSampleableRegion.h"
#include "CUDARRTConnect.h"
#include "CUDAPoissonRRTKernels.h"
#include "CPUUtility.h"
#include "gPlannerCollisionWrapper.h"
#include "GPUMemoryManager.h"
#include "cuda_defs.h"
#include "PlanningParams.h"
#include <cuda_runtime_api.h>
#include "cuda_errorcheck.h"

//#define DEBUG_STEP_VALUES

using namespace ompl;
using namespace poisson;

CUDARRTConnect::CUDARRTConnect(const ompl::base::SpaceInformationPtr &si, const ompl::app::GeometrySpecification* geometrySpecification)
: GPUPlanner(si, geometrySpecification), resetRequired_(false), 
stateValidityCheckingResolutionR3_(0.01f), stateValidityCheckingResolutionSO3_(0.01f)
{
	d_sampleData_.samples = NULL;
	d_sampleData_.isTreeNode = NULL;
	d_sampleData_.levels = NULL;
}

CUDARRTConnect::~CUDARRTConnect()
{
	freeMemory();
}

ompl::base::PlannerStatus CUDARRTConnect::solve(const ompl::base::PlannerTerminationCondition &ptc)
{
	printf("maxdist : %f\n", maxDistance_);
	ompl::time::point timeStart = time::now();

	resetRequired_ = true;

	SerializedStates h_StartAndGoalStates[2];

	// set start state
	const base::State *st = pis_.nextStart();
	if (st == NULL)
	{
		OMPL_ERROR("There are no valid initial states!");
		return base::PlannerStatus::INVALID_START;
	}
	ConvertStateToFloatVector(si_, st, &(h_StartAndGoalStates[0].x[0]));

	// set goal state
	base::Goal                 *goal   = pdef_->getGoal().get();
	base::GoalState			   *goalState = dynamic_cast<base::GoalState*>(goal);
	ConvertStateToFloatVector(si_, goalState->getState(), &(h_StartAndGoalStates[1].x[0]));

	rrtconnect::cuda::setStartAndGoalStates(h_StartAndGoalStates, d_sampleData_, d_startTree_, d_goalTree_);

#ifdef DEBUG_STEP_VALUES
	float randSample[7];
	int nodeIndex;
	float nearestNode[7];
	float destState[7];
	float randSample2[7];
	int colResult;
	float nearestNode2[7];
	float destState2[7];
	int colResult2;
#endif

	bool terminate = false; 
	bool expandStartTree = true;
	std::vector<SerializedStates> solution;
	int iterations = 0;
	do
	{
		RrtTree& currentTree = expandStartTree ? d_startTree_ : d_goalTree_;
		RrtTree& otherTree = expandStartTree ? d_goalTree_ : d_startTree_;

		// current tree
		////////////////////////////////////////////////////////////////////////
		// generate random configurations
		GpuMemoryManager::MemoryChunk* destStatesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(SerializedStates) * NUM_THREADS);
		generateRandomStates((SerializedStates*)destStatesChunk->p);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(randSample, destStatesChunk->p, sizeof(float) * 7, cudaMemcpyDeviceToHost);
#endif
		// compute nearest nodes
		GpuMemoryManager::MemoryChunk* distancesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(float) * NUM_THREADS);
		GpuMemoryManager::MemoryChunk* nodeIndicesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(int) * NUM_THREADS);
		computeNearestNodes(currentTree, (SerializedStates*)destStatesChunk->p, (float*)distancesChunk->p, (int*)nodeIndicesChunk->p);
		GpuMemoryManager::getInstance()->releaseChunk(distancesChunk);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(&nodeIndex, nodeIndicesChunk->p, sizeof(int), cudaMemcpyDeviceToHost);
		cudaMemcpy(nearestNode, currentTree.treeNodes + nodeIndex, sizeof(float) * 7, cudaMemcpyDeviceToHost);
#endif

		GpuMemoryManager::MemoryChunk* connectionCheckChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(int) * NUM_THREADS);
		// change destStates if it is longer than maxDistance
		// connectionCheck is 1 for all query
		setDestStates(currentTree, (int*)nodeIndicesChunk->p, (SerializedStates*)destStatesChunk->p,
			(int*)connectionCheckChunk->p, true);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(destState, destStatesChunk->p, sizeof(float) * 7, cudaMemcpyDeviceToHost);
#endif

		// collision check
		// connectionCheck is set to 0 if collision detected
		GpuMemoryManager::MemoryChunk* resultsChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(int) * NUM_THREADS);
		GpuMemoryManager::MemoryChunk* colFreeStatesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(SerializedStates) * NUM_THREADS);
		checkCollision(currentTree, (int*)nodeIndicesChunk->p, (SerializedStates*)destStatesChunk->p, (int*)resultsChunk->p, 
			(SerializedStates*)colFreeStatesChunk->p, (int*)connectionCheckChunk->p);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(&colResult, resultsChunk->p, sizeof(int), cudaMemcpyDeviceToHost);
#endif

		// add new nodes
		addNewNodes(currentTree, otherTree, (int*)nodeIndicesChunk->p, (SerializedStates*)destStatesChunk->p, (int*)resultsChunk->p, 
			(SerializedStates*)colFreeStatesChunk->p);

		GpuMemoryManager::getInstance()->releaseChunk(resultsChunk);
		GpuMemoryManager::getInstance()->releaseChunk(colFreeStatesChunk);

		// other tree
		////////////////////////////////////////////////////////////////////////

		// update dest states
		updateDestStates((SerializedStates*)destStatesChunk->p, (int*)resultsChunk->p, (SerializedStates*)colFreeStatesChunk->p);

		// compute nearest nodes
		distancesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(float) * NUM_THREADS);
		GpuMemoryManager::MemoryChunk* nodeIndicesChunk2 = GpuMemoryManager::getInstance()->getChunk(sizeof(int) * NUM_THREADS);
		computeNearestNodes(otherTree, (SerializedStates*)destStatesChunk->p, (float*)distancesChunk->p, (int*)nodeIndicesChunk2->p);
		GpuMemoryManager::getInstance()->releaseChunk(distancesChunk);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(&nodeIndex, nodeIndicesChunk2->p, sizeof(int), cudaMemcpyDeviceToHost);
		cudaMemcpy(nearestNode2, otherTree.treeNodes + nodeIndex, sizeof(float) * 7, cudaMemcpyDeviceToHost);
#endif

		// change destStates if it is longer than maxDistance
		// connectionCheck is set to 0 for truncated queries
		setDestStates(otherTree, (int*)nodeIndicesChunk2->p, (SerializedStates*)destStatesChunk->p,
			(int*)connectionCheckChunk->p, false);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(destState2, destStatesChunk->p, sizeof(float) * 7, cudaMemcpyDeviceToHost);
#endif

		// collision check
		// connectionCheck is set to 0 if collision detected
		resultsChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(int) * NUM_THREADS);
		colFreeStatesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(SerializedStates) * NUM_THREADS);
		checkCollision(otherTree, (int*)nodeIndicesChunk2->p, (SerializedStates*)destStatesChunk->p, (int*)resultsChunk->p, 
			(SerializedStates*)colFreeStatesChunk->p, (int*)connectionCheckChunk->p);
#ifdef DEBUG_STEP_VALUES
		cudaMemcpy(&colResult2, resultsChunk->p, sizeof(int), cudaMemcpyDeviceToHost);
#endif

		// add new nodes
		addNewNodes(otherTree, currentTree, (int*)nodeIndicesChunk2->p, (SerializedStates*)destStatesChunk->p, (int*)resultsChunk->p, 
			(SerializedStates*)colFreeStatesChunk->p);
		
		GpuMemoryManager::getInstance()->releaseChunk(resultsChunk);
		GpuMemoryManager::getInstance()->releaseChunk(colFreeStatesChunk);

#ifdef DEBUG_STEP_VALUES
		printf("(%3.2f,%3.2f,%3.2f) %d (%3.2f,%3.2f,%3.2f|%3.2f,%3.2f,%3.2f,%3.2f)\n",
			randSample[0], randSample[1], randSample[2], expandStartTree ? 1 : 2,
			nearestNode[0], nearestNode[1], nearestNode[2], nearestNode[3], nearestNode[4], nearestNode[5], nearestNode[6]);
		printf("(%3.2f,%3.2f,%3.2f) %d (%3.2f,%3.2f,%3.2f|%3.2f,%3.2f,%3.2f,%3.2f) %s\n",
			randSample[0], randSample[1], randSample[2], expandStartTree ? 1 : 2,
			destState[0], destState[1], destState[2], destState[3], destState[4], destState[5], destState[6],
			(colResult == 0) ? "O" : "X");

		printf("(%3.2f,%3.2f,%3.2f) %d (%3.2f,%3.2f,%3.2f|%3.2f,%3.2f,%3.2f,%3.2f)\n",
			randSample2[0], randSample2[1], randSample2[2], expandStartTree ? 2 : 1,
			nearestNode2[0], nearestNode2[1], nearestNode2[2], nearestNode2[3], nearestNode2[4], nearestNode2[5], nearestNode2[6]);
		printf("(%3.2f,%3.2f,%3.2f) %d (%3.2f,%3.2f,%3.2f|%3.2f,%3.2f,%3.2f,%3.2f) %s\n\n",
			randSample2[0], randSample2[1], randSample2[2], expandStartTree ? 2 : 1,
			destState2[0], destState2[1], destState2[2], destState2[3], destState2[4], destState2[5], destState2[6],
			(colResult2 == 0) ? "O" : "X");
#endif

		// termination check
		// if one connectionCheck is set to 1, the solution is found
		terminate = isSolutionFound(currentTree, otherTree, solution, (int*)connectionCheckChunk->p, 
			(int*)nodeIndicesChunk->p, (int*)nodeIndicesChunk2->p, (SerializedStates*)destStatesChunk->p);

		GpuMemoryManager::getInstance()->releaseChunk(nodeIndicesChunk);
		GpuMemoryManager::getInstance()->releaseChunk(nodeIndicesChunk2);
		GpuMemoryManager::getInstance()->releaseChunk(destStatesChunk);
		GpuMemoryManager::getInstance()->releaseChunk(connectionCheckChunk);

		expandStartTree = !expandStartTree;
		++iterations;
	} while (!terminate);

	// last currentTree = goal tree
	if (expandStartTree)
		std::reverse(solution.begin(), solution.end());

	double planningTime = time::seconds(time::now() - timeStart);

	printf("Solution after %d iterations in %f ms\n", iterations, planningTime * 1000);
	for (unsigned int i = 0; i < solution.size(); ++i)
	{
		float euler[3];
		float qx = solution[i].components.q[0];
		float qy = solution[i].components.q[1];
		float qz = solution[i].components.q[2];
		float qw = solution[i].components.q[3];
		euler[0] = atan2f(2*(qx*qw+qy*qz),1-2*(qx*qx+qy*qy));
		euler[1] = asinf(2*(qw*qy-qx*qz));
		euler[2] = atan2f(2*(qz*qw+qx*qy),1-2*(qy*qy+qz*qz));
		euler[0] *= 180.f / 3.141592f;
		euler[1] *= 180.f / 3.141592f;
		euler[2] *= 180.f / 3.141592f;

		float dist = 0;
		if (i > 0)
		{
			dist = (solution[i].x[0]-solution[i-1].x[0]) * (solution[i].x[0]-solution[i-1].x[0]) +
				(solution[i].x[1]-solution[i-1].x[1]) * (solution[i].x[1]-solution[i-1].x[1]) +
				(solution[i].x[2]-solution[i-1].x[2]) * (solution[i].x[2]-solution[i-1].x[2]);
			dist = sqrtf(dist);
		}

		printf("%d (%.3f %.3f %.3f | %.3f %.3f %.3f) adv:%f\n",
			i,
			solution[i].x[0],
			solution[i].x[1],
			solution[i].x[2],
			euler[0], euler[1], euler[2],
			dist);
	}
	printf("\n");
	printf("# of nodes in start tree : %d\n", *d_startTree_.h_numTreeNodes);
	printf("# of nodes in goal tree : %d\n", *d_goalTree_.h_numTreeNodes);
	printf("maxdist : %f\n", maxDistance_);

	bool solved = false;
	bool approximate = false;
	return base::PlannerStatus(solved, approximate);
}

void CUDARRTConnect::clear(void)
{
	GPUPlanner::clear();

	if (resetRequired_)
		reset();
}

void CUDARRTConnect::reset()
{
	// TODO:
	// set numSmples as numPrecomputedSamples
	// reset trees

	resetRequired_ = false;
}

void CUDARRTConnect::setup(void)
{
	GPUPlanner::setup();

	// bounds
	std::vector<float> vBoundsMin(3);
	std::vector<float> vBoundsMax(3);
	for (int i = 0; i < 3; ++i)
	{
		vBoundsMin[i] = (float)GetBoundsMin()[i];
		vBoundsMax[i] = (float)GetBoundsMax()[i]; 
	}

	// state validity checking resolution
	base::StateSpacePtr spaceR3 = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(0);
	base::StateSpacePtr spaceSO3 = si_->getStateSpace()->as<base::CompoundStateSpace>()->getSubspace(1);
	stateValidityCheckingResolutionR3_ = (float)((double)spaceR3->getValidSegmentCountFactor()
		/ (spaceR3->getMaximumExtent() * spaceR3->getLongestValidSegmentFraction()));
	stateValidityCheckingResolutionSO3_ = (float)((double)spaceSO3->getValidSegmentCountFactor()
		/ (spaceSO3->getMaximumExtent() * spaceSO3->getLongestValidSegmentFraction()));

	// meshes
	GMesh* h_obstacle;
	GMesh* h_robot;
	gplanner_collision_wrapper::setup(geometrySpecification_, &h_obstacle, &h_robot);

	// one-time memory allocation
	initializeMemory();

	// upload
	uploadData(vBoundsMin, vBoundsMax, h_robot, h_obstacle);

	// initialize random states
	poisson::cuda::initializeRandomStates();

	GpuMemoryManager::getInstance()->init();
}

void CUDARRTConnect::freeMemory(void)
{
	GPUPlanner::freeMemory();

	rrtconnect::cuda::freeMemory(d_sampleData_, d_robot_, d_obstacle_, d_startTree_, d_goalTree_);
	
	GpuMemoryManager::getInstance()->destroy();
}

void CUDARRTConnect::initializeMemory()
{
	poisson::cuda::initializeMemory(d_sampleData_, &d_robot_, &d_obstacle_, d_startTree_, d_goalTree_);
}

void CUDARRTConnect::uploadData(const std::vector<float>& vBoundsMin, const std::vector<float>& vBoundsMax, 
	const GMesh* h_robot, const GMesh* h_obstacle)
{
	rrtconnect::cuda::uploadData(vBoundsMin, vBoundsMax, h_robot, h_obstacle, 
		d_robot_, d_obstacle_, d_sampleData_);
}

// primitive methods
////////////////////////////////////////////////////////////////////////////////
void CUDARRTConnect::generateRandomStates(SerializedStates* outStates)
{
	poisson::cuda::generateRandomStates(outStates);
}

void CUDARRTConnect::computeNearestNodes(RrtTree& tree, SerializedStates* inStates, float* outDistances, int* outNodeIndices)
{
	int numTreeNodes = (int)(*(tree.h_numTreeNodes));
	GpuMemoryManager::MemoryChunk* internalDistancesChunk = GpuMemoryManager::getInstance()->getChunk(sizeof(float) * NUM_THREADS * numTreeNodes);
	float* internalDistances = (float*)internalDistancesChunk->p;

	poisson::cuda::knn(tree.treeNodes, numTreeNodes, inStates, NUM_THREADS, 1, internalDistances, outDistances, outNodeIndices);

	GpuMemoryManager::getInstance()->releaseChunk(internalDistancesChunk);
}

void CUDARRTConnect::setDestStates(RrtTree& tree, int* inNodeIndices, SerializedStates* inoutDestStates,
	int* outConnectionCheck, bool isCurrentTree)
{
	rrtconnect::cuda::setDestStates(NUM_THREADS, tree.treeNodes, inNodeIndices, inoutDestStates, outConnectionCheck, isCurrentTree, (float)maxDistance_);
}

void CUDARRTConnect::checkCollision(RrtTree& currentTree, int* inNodeIndices, SerializedStates* inDestStates, 
	int* outResults, SerializedStates* outColFreeStates, int* inoutConnectionCheck)
{
	rrtconnect::cuda::checkCollision(NUM_THREADS, currentTree.treeNodes, inNodeIndices, inDestStates,
		outResults, outColFreeStates,
		d_robot_, d_obstacle_, stateValidityCheckingResolutionR3_, stateValidityCheckingResolutionSO3_, inoutConnectionCheck);
}

void CUDARRTConnect::addNewNodes(RrtTree& currentTree, RrtTree& otherTree, int* inNodeIndices, SerializedStates* inDestStates, 
	int* inResults, SerializedStates* inColFreeStates)
{
	rrtconnect::cuda::addNewNodes(NUM_THREADS, currentTree, otherTree, inNodeIndices, inDestStates, inResults, inColFreeStates);

	rrtconnect::cuda::downloadValues(currentTree);
}

void CUDARRTConnect::updateDestStates(SerializedStates* inoutDestStates, int* inResults, SerializedStates* inColFreeStates)
{
	rrtconnect::cuda::updateDestStates(NUM_THREADS, inoutDestStates, inResults, inColFreeStates);
}

bool CUDARRTConnect::isSolutionFound(RrtTree& currentTree, RrtTree& otherTree, std::vector<SerializedStates>& solution, 
	int* inConnectionCheck, int* inCurrentTreeNodeIndices, int* inOtherTreeNodeIndices, SerializedStates* inDestStates)
{
	return rrtconnect::cuda::isSolutionFound(NUM_THREADS, currentTree, otherTree, solution, inConnectionCheck, inCurrentTreeNodeIndices, 
		inOtherTreeNodeIndices, inDestStates);
}

